import os
import pickle
from typing import List

from langchain_core.documents import Document
from langchain_text_splitters import (
    MarkdownHeaderTextSplitter,
    RecursiveCharacterTextSplitter,
)

from configs import app_config
from logger import logger

# 加载政策文件进行
class DocumentProcessor:
    def __init__(
        self,
        md_directory=app_config.DOCUMENT_MD_DIR,
        chunk_size=app_config.CHUNK_SIZE,
        chunk_overlap=app_config.OVERLAP_SIZE,
    ):
        """
        初始化DocumentProcessor类
        
        参数:
        md_directory: Markdown文件所在的目录
        chunk_size: 文档分割的大小
        chunk_overlap: 文档分割的重叠部分大小
        """
        self.md_directory = md_directory
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def __get_documents_from_md(self, md_path):
        """
        处理单个markdown文件，生成文档片段
        
        参数:
        md_path: Markdown文件的路径
        
        返回:
        List[Document]: 处理后的文档片段列表
        """
        logger.info(f"开始处理文档: {md_path}")
        with open(md_path, "r") as f:
            md_content = f.read()

        # 定义Markdown标题层级
        headers_to_split_on = [
            ("#", "h1"),
            ("##", "h2"),
            ("###", "h3"),
        ]

        # 使用MarkdownHeaderTextSplitter进行初步分割，即首先按照一级标题或者二级标题进行分割
        """
        例如：
        片段1
           # 第一章
           这是第一章的内容。
        片段2
           ## 第二章
           这是第二章的内容。
        片段3
           ## 第三章
           这是第三章的内容。
        """
        md_splitter = MarkdownHeaderTextSplitter(headers_to_split_on)
        md_header_splits = md_splitter.split_text(md_content)

        # 使用RecursiveCharacterTextSplitter进行进一步分割，chunk_size为300，overlap为30
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            keep_separator="end",
        )
        documents: List[Document] = text_splitter.split_documents(md_header_splits)

        # 对每个文档片段进行后处理
        documents = list(
            map(self.__post_process_document, documents, [md_path] * len(documents))
        )
        logger.info(f"文档处理完成,共生成 {len(documents)} 个文档片段")
        return documents

    def __post_process_document(self, document: Document, md_path: str):
        """
        对单个文档片段进行后处理
        
        参数:
        document: 待处理的文档片段
        md_path: 原始Markdown文件的路径
        
        返回:
        Document: 处理后的文档片段
        """
        # 例如：md_path = "CrossBorderAIChallenge/data/policy/prohibited/prohibited_1.md"
        file_name = md_path.split("/")[-1]
        header = ""
        # 根据文件名判断文档类型，是禁售政策文件还是合规政策文件
        if "prohibited" in file_name:
            header += "prohibited | "
        else:
            header += "requires_compliance_check | "
        document.metadata.update({"file_name": file_name})
        # 确保h1存在，如果没有则设为空字符串
        h1 = document.metadata.get("h1", "")
        if not h1:
            document.metadata["h1"] = ""
        # 构建header字符串,包含除file_name外的所有metadata
        metadata_values = []
        for key, value in document.metadata.items():
            if key != "file_name":
                metadata_values.append(value)
        header += " | ".join(metadata_values)
        document.page_content = document.page_content.strip()
        # 将header信息添加到文档内容的开头
        document.page_content = f"header:{header}\n content: {document.page_content}"
        return document
    # 获取指定目录下所有的Markdown文件路径
    def __get_all_markdown_files(self) -> List[str]:
        """
        获取指定目录下所有的Markdown文件路径
        
        返回:
        List[str]: Markdown文件路径列表
        """
        markdown_files = []
        for root, dirs, files in os.walk(self.md_directory):
            for file in files:
                if file.endswith(".md"):
                    markdown_files.append(os.path.join(root, file))
        return markdown_files

    def get_all_documents(self, cache=False) -> List[Document]:
        """
        获取所有文档
        
        参数:
        cache: 是否使用缓存（此参数在当前方法中未使用，可能是为了与get_document_pipeline方法保持一致）
        
        返回:
        List[Document]: 所有处理后的文档片段列表
        """
        logger.info("开始获取所有文档")
        all_md_files = self.__get_all_markdown_files()

        documents = []
        # 循环处理每个Markdown文件，将每个md文件处理后得到的文档片段放在一个列表中
        for md_file in all_md_files:
            documents.extend(self.__get_documents_from_md(md_file))
        logger.info(f"获取完成,共 {len(documents)} 个文档")
        return documents

    def get_document_pipeline(self, cache=True) -> List[Document]:
        """
        获取文档处理流水线，支持缓存
        
        参数:
        cache: 是否使用缓存
        
        返回:
        List[Document]: 处理后的文档片段列表
        """
        cache_path = "cache/documents.pkl"

        if cache and os.path.exists(cache_path):
            # 如果启用缓存且缓存文件存在，则从缓存加载文档
            with open(cache_path, "rb") as f:
                documents = pickle.load(f)
            logger.info(f"从缓存加载文档，共 {len(documents)} 个")
        else:
            # 否则重新处理所有文档
            documents = self.get_all_documents()
            # 将处理结果保存到缓存
            with open(cache_path, "wb") as f:
                pickle.dump(documents, f)
            logger.info(f"文档已保存到缓存 {cache_path}")

        return documents

if __name__ == "__main__":
    processor = DocumentProcessor()
    processor.get_document_pipeline()
